import sys
import os
from nltk.translate.bleu_score import corpus_bleu
import codecs
from collections import OrderedDict

discourse_sympol = {'<D%s>' % i for i in range(8)}

MAX_SEN_LEN = 250


def contain_sep(item):
    return ' <SEP> ' in item


def clear_text(text):
    text = text.strip()
    global is_filter
    if is_filter == 'true':
        text = text.split(' <SEP> ')[1]
    for dis in discourse_sympol:
        text = text.replace(' ' + dis, '')
    text = text.split(' ')[:MAX_SEN_LEN]
    text = ' '.join(text)
    if not text.endswith('.'):
        last_index = text.rfind('.')
        if last_index != -1:
            text = text[:last_index + 1]
    return text


def convert_hyp(text):
    text = clear_text(text)
    return text.split(' ')


def convert_ref(text):
    text = clear_text(text)
    return [text.split(' ')]


def calc_diversity(texts):
    unigram = set()
    bigram = set()
    num_tok = 0
    for vec in texts:
        v_len = len(vec)
        num_tok += v_len
        unigram.update(vec)
        bigram.update([tuple(vec[i:i+2]) for i in range(v_len-1)])
    metrics = OrderedDict()
    metrics['d_1'] = round(len(unigram) * 1.0 / num_tok * 100, 6)
    metrics['d_2'] = round(len(bigram) * 1.0 / num_tok * 100, 6)
    metrics['num_d1'] = len(unigram)
    metrics['num_d2'] = len(bigram)
    metrics['num_tok'] = num_tok
    metrics['sen_len'] = round(num_tok * 1.0 / len(texts), 6)
    return metrics


hyp_path = sys.argv[1]
is_filter = sys.argv[2]


hypos = codecs.open(hyp_path, 'r', encoding='utf8').readlines()
print('total number of test example', len(hypos))

if is_filter == 'true':
    hypos = filter(contain_sep, hypos)

hypos = map(convert_hyp, hypos)
print('correct number of test example', len(hypos))
print('hyp')
print(' '.join(hypos[0]))

metrics = calc_diversity(hypos)

print_distinct = 'd1={:.6f}, d2={:.6f}, nd1={:d}, nd2={:d}, ntok={:d}, avg_sen_len={:.6f}'.format(*metrics.values())
print('=' * 50)
print(print_distinct)
print('=' * 50)

# write to file
with open(os.path.join(hyp_path + '.distinct'), 'w') as f:
    f.write(print_distinct)



